package org.example.service;

import io.grpc.stub.StreamObserver;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

/**
 * @Author wangjun
 * @Date 2025/3/20 16:01
 */
public class GrpcRetryUtil {
    public static <T, R> void executeWithRetry(
            Supplier<T> stubSupplier,  // 存根生成器
            GrpcRetryOperation<T, R> operation,  // 业务逻辑
            int maxRetries,  // 最大重试次数
            long initialBackoffMs,  // 初始退避时间
            Class<? extends R> responseType  // 响应类型（用于类型安全）
    ) throws Exception {
        int retryCount = 0;
        Throwable lastError = null;

        while (retryCount <= maxRetries) {
            final CountDownLatch latch = new CountDownLatch(1);
            final AtomicBoolean success = new AtomicBoolean(false);
            final AtomicReference<Throwable> errorRef = new AtomicReference<>();

            try {
                T stub = stubSupplier.get();
                StreamObserver<R> responseObserver = new StreamObserver<R>() {
                    @Override
                    public void onNext(R value) {
                        System.out.println("收到响应: " + value);
                    }

                    @Override
                    public void onError(Throwable t) {
                        errorRef.set(t);
                        latch.countDown();
                    }

                    @Override
                    public void onCompleted() {
                        success.set(true);
                        latch.countDown();
                    }
                };

                // 执行具体业务逻辑
                operation.execute(stub, responseObserver);

                // 设置合理的超时时间
                if (!latch.await(5, TimeUnit.SECONDS)) {
                    throw new TimeoutException("gRPC调用超时");
                }

                if (success.get()) {
                    return;
                }
                lastError = errorRef.get();
            } catch (Exception e) {
                lastError = e;
            }
            if (retryCount++ >= maxRetries) break;
            // 指数退避计算
            long backoffMs = initialBackoffMs * (1 << (retryCount - 1));
            System.out.printf("第%d次重试，%.2fs后重试...%n", retryCount, backoffMs / 1000.0);
            Thread.sleep(backoffMs);
        }
        throw new Exception("重试失败（" + maxRetries + "次）", lastError);
    }
}

